package com.demo.modules.test.user.utils;

import com.demo.common.redis.service.RedisService;
import com.demo.modules.test.user.mapper.CommonMapper;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.shardingsphere.core.rule.ShardingRule;
import org.apache.shardingsphere.core.rule.TableRule;
import org.apache.shardingsphere.shardingjdbc.jdbc.core.datasource.ShardingDataSource;
import org.apache.shardingsphere.underlying.common.rule.DataNode;
import org.springframework.boot.CommandLineRunner;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import javax.sql.DataSource;
import java.lang.reflect.Field;
import java.lang.reflect.Modifier;
import java.time.LocalDateTime;
import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 项目启动后 读取已有分表 进行缓存
 * @author molong
 * @since 2024/6/26
 */
@Slf4j
@Order(value = 1) // 数字越小 越先执行
@Component
public class ShardingTablesLoadRunner implements CommandLineRunner {

    @Resource
    private CommonMapper commonMapper;

    @Resource
    private DataSource dataSource;

    @Resource
    private RedisService redisService;


    @Override
    public void run(String... args) throws Exception {

        // 给 分表工具类注入属性
        DateShardingAlgorithm.setRedisService(redisService);
        DateShardingAlgorithm.setCommonMapper(commonMapper);

        // 读取数据库中所有表名
        List<String> tableNameList = commonMapper.getAllTableNameBySchema("test");
        // 删除旧的缓存(如果存在)
        redisService.del(RandomUserUtil.TABLE_NAME_CACHE);
        if(CollectionUtils.isNotEmpty(tableNameList)){
            // 写入新的缓存
            redisService.sAdd(RandomUserUtil.TABLE_NAME_CACHE, tableNameList.toArray(new String[0]));
        }
        //更新分表节点配置
        //sharding数据源
        ShardingDataSource shardingDataSource = (ShardingDataSource) this.dataSource;
        //所有的拆分表
        ShardingRule rule = shardingDataSource.getRuntimeContext().getRule();
        Collection<TableRule> tableRules = rule.getTableRules();
        for (TableRule tableRule : tableRules) {
            // 表的节点
            List<DataNode> dataNodes = tableRule.getActualDataNodes();
            // 动态刷新actualDataNodesField
            Field actualDataNodesField = TableRule.class.getDeclaredField("actualDataNodes");
            Field modifiersField = Field.class.getDeclaredField("modifiers");
            modifiersField.setAccessible(true);
            // 设置修饰符：private
            modifiersField.setInt(actualDataNodesField, actualDataNodesField.getModifiers() & ~Modifier.FINAL);

            // 新节点
            List<DataNode> newDataNodes = new ArrayList<>();
            // 数据源
            String dataSourceName = dataNodes.get(0).getDataSourceName();
            // 逻辑表名
            String logicTableName = tableRule.getLogicTable();
            // 循环动态拼接所有节点表名
            String prefix=logicTableName+"_";
            for (String tableName:tableNameList) {
                if(tableName.contains(prefix)){
                    DataNode dataNode = new DataNode(dataSourceName+"."+tableName);
                    newDataNodes.add(dataNode);
                }
            }
            actualDataNodesField.setAccessible(true);
            // 数据更新回去
            actualDataNodesField.set(tableRule, newDataNodes);
            Set<String> actualTables = Sets.newHashSet();
            Map<DataNode, Integer> dataNodeIntegerMap = Maps.newHashMap();

            AtomicInteger a = new AtomicInteger(0);
            newDataNodes.forEach((dataNode -> {
                actualTables.add(dataNode.getTableName());
                if (a.intValue() == 0) {
                    a.incrementAndGet();
                    dataNodeIntegerMap.put(dataNode, 0);
                }
                else {
                    dataNodeIntegerMap.put(dataNode, a.intValue());
                    a.incrementAndGet();
                }
            }));

            // 动态刷新：actualTables
            Field actualTablesField = TableRule.class.getDeclaredField("actualTables");
            actualTablesField.setAccessible(true);
            actualTablesField.set(tableRule, actualTables);

            // 动态刷新：dataNodeIndexMap
            Field dataNodeIndexMapField = TableRule.class.getDeclaredField("dataNodeIndexMap");
            dataNodeIndexMapField.setAccessible(true);
            dataNodeIndexMapField.set(tableRule, dataNodeIntegerMap);

            // 动态刷新：datasourceToTablesMap
            Map<String, Collection<String>> datasourceToTablesMap = Maps.newHashMap();
            datasourceToTablesMap.put(dataSourceName, actualTables);
            Field datasourceToTablesMapField = TableRule.class.getDeclaredField("datasourceToTablesMap");
            datasourceToTablesMapField.setAccessible(true);
            datasourceToTablesMapField.set(tableRule, datasourceToTablesMap);
        }
        log.info("ShardingTablesLoadRunner start OK");
    }
}

